import copy
import json
import os

from baseconf import BASE_DISK

with open("./autofrist_result_pre.json", 'r') as f:
    parsed_data = json.load(f)
from transformers import pipeline

#     single_video_result = {
#         "clip_id": clip_id,
#         "scerario": "cityroad",
#         "weather": "unknown",
#         "period": "night",
#         "road_structure": "ramp",
#         "general_obstacle": "nothing",
#         "abnormal_condition": "nothing",
#         "ego_car_behavior": "turning right",
#         "closest_participants_type": "passenger car",
#         "closest_participants_behavior": "braking"
#     }




model_path="/model_path/mdeberta-v3-base-squad2" if os.path.exists(
                        "/etc") else BASE_DISK + ":/model_path/mdeberta-v3-base-squad2"
qa_model = pipeline(task="question-answering", model=model_path,device="cuda")

submit_json_bert = copy.deepcopy(parsed_data)
for single_result in submit_json_bert["test_results"]:
    for key in single_result.keys():
        if key in ["clip_id", "scerario", "weather", "period", "road_structure"]:
            continue
        desc = single_result[key]

        question = f"What's the {key} like?"
        context = desc
        answer=qa_model(question=question, context=context)['answer']
        single_result[key] = answer
        print(f"{key}:{single_result[key]},{desc}")


with open(file="./autofrist_result3.json",encoding="utf-8",mode="w") as f:
    json_data = json.dumps(submit_json_bert)
    f.write(json_data)